#%%
from torch import nn
import torch
from flask import Flask, jsonify, request, render_template
from flask_cors import CORS
import numpy as np
from flask import Flask, jsonify, request, render_template
import json
app = Flask(__name__)
# select device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)



# ----------------------------------------------------#
#   创建模型
# ----------------------------------------------------#
class TimeSeriesClassifier(nn.Module):
    def __init__(self, n_features, hidden_dim=64, output_size=1):
        super().__init__()
        self.lstm = nn.LSTM(input_size=n_features, hidden_size=hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_size)  # output_size classes

    def forward(self, x):
        x, _ = self.lstm(x)  # LSTM层
        x = x[:, -1, :]  # 只取LSTM输出中的最后一个时间步
        x = self.fc(x)  # 通过一个全连接层
        return x


# ----------------------------------------------------#
#   模型实例化
# ----------------------------------------------------#
#seq_len = 100  # 根据你的序列长度进行调整
n_features = 14  # 根据你的特征数量进行调整
output_size = 4
model = TimeSeriesClassifier(n_features=n_features, output_size=output_size)
#model = torch.load('LSTM_best_model.pth')
checkpoint = torch.load('LSTM_best_model.pth', weights_only=False)
model.load_state_dict(checkpoint['state_dict']) 
model.eval()
model.to(device)
#%%
@app.route('/predict', methods=['POST'])
@torch.no_grad()
def predict():
    if request.method == 'POST':
        file = request.files['file']
        print (file)
        data = np.load(file)
        data = data.reshape(1, 50, 14)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        data = torch.tensor(data).float().to(device)
        # 生成预测并获取最可能的类别
        outputs = model(data)
        _, predicted = torch.max(outputs, 1)
        results = predicted.cpu()
        results = results.numpy()
        dic={}
        dic['index']=results.tolist()
        dicJson = json.dumps(dic)
        # 此处尝试将tensor直接转化为json格式，无法完成，需要在上一层将标签直接变成json格式2024/07/21 
        # 应该就差一步了！！！！！！
        # 20240722卧槽成了！！！！
        return jsonify(dicJson)

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=5000)






# weights_path = ".LSTM_best_model.pth"

# model = model.to(device)
# model.eval()


# class TimeSeriesClassifier(nn.Module):
#     def __init__(self, n_features, hidden_dim=256, output_size=1):
#         super().__init__()
#         self.lstm = nn.LSTM(input_size=n_features, hidden_size=hidden_dim, batch_first=True)
#         self.fc = nn.Linear(hidden_dim, output_size)  # output_size classes

#     def forward(self, x):
#         x, _ = self.lstm(x)  # LSTM层
#         x = x[:, -1, :]  # 只取LSTM输出中的最后一个时间步
#         x = self.fc(x)  # 通过一个全连接层
#         return x


# # ----------------------------------------------------#
# #   模型实例化
# # ----------------------------------------------------#
# # seq_len = 50  # 根据你的序列长度进行调整
# n_features = 14  # 根据你的特征数量进行调整
# output_size = 4
# model = TimeSeriesClassifier(n_features=n_features, output_size=output_size)
# print(model)
# model.load_state_dict(torch.load(weights_path, map_location=device))